@mxnet-label-bot Add [Backend, Operator]

@anirudh2290 I've simplified the exception catching according to your suggestion.

Agreed that the situation 2 is really tricky, because for ExecType::kAsync op, the on_complete callback is async to op computation, and there are no mechanism to ensure on_complete not skipped when exception happens.

Now ExecType::kAsync is only used in custom op.

@anirudh2290 I'll test the example you've mentioned.

Let me mention another problem here:
Adding WaitToRead in custom thread will block the thread, which will cause deadlock. Suppose we only have one custom thread and custom op A is dependent upon custom op B (A->B), then deadlock will happen due to not enough working custom thread to use for B.

import mxnet as mx

class A(mx.operator.CustomOp):
    def __init__(self):
        super(A, self).__init__()
    def forward(self, is_train, req, in_data, out_data, aux):
        out_data[0][:] = mx.nd.Custom(in_data[0], in_data[1], op_type='B')
    def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
        in_grad[0][:] = out_grad[0]
        in_grad[1][:] = out_grad[0]

@mx.operator.register("A")
class AProp(mx.operator.CustomOpProp):
    def __init__(self):
        super(AProp, self).__init__()
    def list_arguments(self):
        return ['a', 'b']
    def list_outputs(self):
        return ['output']
    def infer_shape(self, in_shape):
        return in_shape, [in_shape[0]]
    def create_operator(self, ctx, shapes, dtypes):
        return A()

class B(mx.operator.CustomOp):
    def __init__(self):
        super(B, self).__init__()
    def forward(self, is_train, req, in_data, out_data, aux):
        out_data[0][:] = in_data[0] + in_data[1]
    def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
        in_grad[0][:] = out_grad[0]
        in_grad[1][:] = out_grad[0]

@mx.operator.register("B")
class BProp(mx.operator.CustomOpProp):
    def __init__(self):
        super(BProp, self).__init__()
    def list_arguments(self):
        return ['a', 'b']
    def list_outputs(self):
        return ['output']
    def infer_shape(self, in_shape):
        return in_shape, [in_shape[0]]
    def create_operator(self, ctx, shapes, dtypes):
        return B()

def main():
    a = mx.nd.array([1, 2])
    b = mx.nd.array([4, 5])
    c = mx.nd.Custom(a, b, op_type='A')
    print c.asnumpy()

if __name__ == '__main__':
    main()

Run the code with:

MXNET_CUSTOM_OP_NUM_THREADS=1 python reproduce.py

Even if you have n threads, we can still deadlock the execution with n+1 custom op dependency chain: A->B->C...

We should make sure there are always enough threads to use, maybe remove MXNET_CUSTOM_OP_NUM_THREADS and create new thread when in need?

@arcadiaphy Custom Operator creates new thread when in need now, and MXNET_CUSTOM_OP_NUM_THREADS is the maximum number of threads.

Removing this line https://github.com/apache/incubator-mxnet/blob/master/src/operator/custom/custom-inl.h#L185 to support unlimited number of threads.

@arcadiaphy Yes good point, the code you linked will cause a deadlock. this is also true if custom op forward or backward has asnumpy() or wait_to_read at the end even before applying your change, when you change MXNET_CUSTOM_OP_NUM_THREADS=1. Even before the custom op multi threaded change this was the case.

Also, I am concerned about the performance because of the blocking call during op execution. I think we have reached a point where there is no option but to change the approach and special case "CustomOperator" in ExecuteOprBlock by checking whether op is CustomOperator and if so, execute the operator instead of skipping (https://github.com/apache/incubator-mxnet/blob/master/src/engine/threaded_engine.h#L377) . I was trying to avoid this to keep the engine agnostic of the op its executing but this seems like the least impactful solution currently.

@anirudh2290 If the engine code remains unchanged, I think enabling unlimited number of custom threads when needed is a quick fix.

Also for unlimited threads, theoretically, there should be no performance impact.

I think we should make the modification in engine. Keeping that env variable is fine as Users may want to limit the custom op threads, so that customers don't run into situation where there are too many custom ops in their model and lot of threads are spawned and thread spawning overhead being high. Also, performance wise adding the waittoread and waittowrite can impact performance: every thread is now pushing all ops, waiting it to complete and then exiting as compared to custom op threads earlier which would just push op to engine and exit. Custom op threads were lightweight earlier and we should keep it that way.

@anirudh2290 How about leave this PR temporarily and create a new PR to modify the engine code?

Yes I am fine with that

close due to issue fixed in #14693